package com.ti.demo.mcp;

import io.modelcontextprotocol.spec.McpSchema;
import org.springaicommunity.mcp.annotation.McpComplete;
import org.springframework.stereotype.Component;

import java.util.*;
import java.util.stream.Collectors;

/**
 * @description: @McpComplete
 * The @McpComplete annotation provides auto-completion functionality for prompts.
 * @Author Triagen
 * @Date 2025/10/30 13:55
 */
@Component
public class McpCompletes {

    // 城市数据列表
    private final List<String> cities = Arrays.asList(
            "北京", "上海", "广州", "深圳", "杭州", "苏州", "南京", "成都", "重庆", "武汉",
            "西安", "天津", "青岛", "厦门", "昆明", "大连", "哈尔滨", "沈阳", "长沙", "宁波",
            "beijing", "shanghai", "guangzhou", "shenzhen", "hangzhou"
    );

    // 国家数据列表
    private final List<String> countries = Arrays.asList(
            "中国", "美国", "日本", "韩国", "泰国", "新加坡", "马来西亚", "越南", "英国", "法国",
            "德国", "意大利", "西班牙", "澳大利亚", "加拿大", "俄罗斯", "巴西", "印度", "chinese",
            "american", "japanese", "korean", "thai", "singapore"
    );

    // 预定义的代码片段和关键字
    private final Map<String, List<String>> codeSnippets = new HashMap<>();
    private final Set<String> javaKeywords = Set.of(
            "public", "private", "protected", "class", "interface", "void", "int",
            "String", "boolean", "if", "else", "for", "while", "return", "new"
    );

    public McpCompletes() {
        initializeCodeSnippets();
    }

    private void initializeCodeSnippets() {
        // 常用方法模板
        codeSnippets.put("public", Arrays.asList(
                "public class ${1:ClassName} {",
                "public void ${1:methodName}() {",
                "public static void main(String[] args) {"
        ));

        codeSnippets.put("private", Arrays.asList(
                "private String ${1:fieldName};",
                "private void ${1:methodName}() {"
        ));

        codeSnippets.put("for", Arrays.asList(
                "for (int i = 0; i < ${1:length}; i++) {",
                "for (String item : ${1:collection}) {"
        ));

        codeSnippets.put("if", Arrays.asList(
                "if (${1:condition}) {",
                "if (${1:condition}) { } else { }"
        ));

        // 常用Spring注解
        codeSnippets.put("@", Arrays.asList(
                "@Service", "@Controller", "@RestController", "@Autowired",
                "@RequestMapping", "@GetMapping", "@PostMapping"
        ));
    }

    // Basic Usage
    @McpComplete(prompt = "city-search")
    public List<String> completeCityName(String prefix) {
        return cities.stream()
                .filter(city -> city.toLowerCase().startsWith(prefix.toLowerCase()))
                .limit(10)
                .toList();
    }


    // With CompleteRequest.CompleteArgument
    @McpComplete(prompt = "travel-planner")
    public List<String> completeTravelDestination(McpSchema.CompleteRequest.CompleteArgument argument) {
        String prefix = argument.value().toLowerCase();
        String argumentName = argument.name();

        // Different completions based on argument name
        if ("city".equals(argumentName)) {
            return completeCities(prefix);
        } else if ("country".equals(argumentName)) {
            return completeCountries(prefix);
        }

        return List.of();
    }

    /**
     * 城市自动补全
     */
    private List<String> completeCities(String prefix) {
        return cities.stream()
                .filter(city -> city.toLowerCase().startsWith(prefix.toLowerCase()))
                .limit(10) // 限制返回数量，提高性能
                .collect(Collectors.toList());
    }

    /**
     * 国家自动补全
     */
    private List<String> completeCountries(String prefix) {
        return countries.stream()
                .filter(country -> country.toLowerCase().startsWith(prefix.toLowerCase()))
                .limit(10)
                .collect(Collectors.toList());
    }


    // With CompleteResult
    @McpComplete(prompt = "code-completion")
    public McpSchema.CompleteResult completeCode(String prefix) {
        List<String> completions = generateCodeCompletions(prefix);
        boolean hasMoreCompletions = completions.size() > 10; // 如果结果多于10个，标记有更多

        // 限制返回数量，提高性能
        if (completions.size() > 10) {
            completions = completions.subList(0, 10);
        }

        return new McpSchema.CompleteResult(
                new McpSchema.CompleteResult.CompleteCompletion(
                        completions,
                        completions.size(),  // total
                        hasMoreCompletions   // hasMore flag
                )
        );
    }

    /**
     * 生成代码补全建议
     */
    private List<String> generateCodeCompletions(String prefix) {
        if (prefix == null || prefix.trim().isEmpty()) {
            return getPopularCompletions();
        }

        String lowerPrefix = prefix.toLowerCase();
        List<String> suggestions = new ArrayList<>();

        // 1. 关键字补全
        suggestions.addAll(javaKeywords.stream()
                .filter(keyword -> keyword.toLowerCase().startsWith(lowerPrefix))
                .toList());

        // 2. 代码片段补全
        suggestions.addAll(codeSnippets.entrySet().stream()
                .filter(entry -> entry.getKey().toLowerCase().startsWith(lowerPrefix))
                .flatMap(entry -> entry.getValue().stream())
                .toList());

        // 3. 智能推断补全（基于常见模式）
        suggestions.addAll(generateIntelligentCompletions(prefix));

        // 去重并排序
        return suggestions.stream()
                .distinct()
                .sorted(Comparator.comparing(String::length)) // 短的建议优先
                .toList();
    }

    /**
     * 基于输入前缀生成智能补全建议
     */
    private List<String> generateIntelligentCompletions(String prefix) {
        List<String> intelligentSuggestions = new ArrayList<>();

        // 根据输入内容推断可能的补全
        if (prefix.startsWith("@")) {
            // Spring注解补全
            intelligentSuggestions.addAll(Arrays.asList(
                    "@Service", "@Component", "@Repository", "@Autowired",
                    "@Value", "@Configuration", "@Bean"
            ));
        } else if (prefix.contains(".")) {
            // 方法调用补全
            String[] parts = prefix.split("\\.");
            if (parts.length > 1) {
                String lastPart = parts[parts.length - 1];
                intelligentSuggestions.addAll(generateMethodCompletions(lastPart));
            }
        } else if (prefix.matches("[a-zA-Z]+")) {
            // 变量名补全
            intelligentSuggestions.addAll(generateVariableCompletions(prefix));
        }

        return intelligentSuggestions.stream()
                .filter(suggestion -> suggestion.toLowerCase().contains(prefix.toLowerCase()))
                .toList();
    }

    /**
     * 生成方法相关的补全建议
     */
    private List<String> generateMethodCompletions(String prefix) {
        List<String> methods = new ArrayList<>();

        if (prefix.startsWith("get")) {
            methods.add("get" + capitalize(prefix.substring(3)) + "()");
        } else if (prefix.startsWith("set")) {
            methods.add("set" + capitalize(prefix.substring(3)) + "(${1:value})");
        } else if (prefix.startsWith("is")) {
            methods.add("is" + capitalize(prefix.substring(2)) + "()");
        }

        // 常用方法模式
        methods.addAll(Arrays.asList(
                "toString()", "equals()", "hashCode()", "length()", "size()",
                "add()", "remove()", "contains()", "isEmpty()"
        ));

        return methods;
    }

    /**
     * 生成变量名补全建议
     */
    private List<String> generateVariableCompletions(String prefix) {
        return Arrays.asList(
                prefix + "List", prefix + "Map", prefix + "Set", prefix + "Array",
                "new" + capitalize(prefix) + "()", prefix + "Service", prefix + "Controller"
        );
    }

    /**
     * 返回常用的补全建议（当输入为空时）
     */
    private List<String> getPopularCompletions() {
        return Arrays.asList(
                "public class ", "public void ", "private String ",
                "List<String> ", "Map<String, Object> ", "@Service", "@Autowired",
                "for (", "if (", "return ", "new "
        );
    }

    /**
     * 字符串首字母大写
     */
    private String capitalize(String str) {
        if (str == null || str.isEmpty()) return str;
        return str.substring(0, 1).toUpperCase() + str.substring(1);
    }

    /**
     * 添加自定义代码片段的方法（可选）
     */
    public void addCodeSnippet(String trigger, String snippet) {
        codeSnippets.computeIfAbsent(trigger, k -> new ArrayList<>()).add(snippet);
    }

    /**
     * 获取所有可用的补全触发器（用于调试）
     */
    public Set<String> getAvailableCompletions() {
        Set<String> allCompletions = new HashSet<>(javaKeywords);
        allCompletions.addAll(codeSnippets.keySet());
        return allCompletions;
    }

}
